package main

import (
	"flag"
	"fmt"
	"github.com/sirupsen/logrus"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

const (
	IPPROTO_OSPF   = 89
	OSPFVersion    = 2
	OSPFHelloType  = 1
	OSPFHeaderLen  = 24
	HelloPacketLen = 20
	snapLen        = 65536
	promiscuous    = false
	timeout        = pcap.BlockForever
)

type OSPFServer struct {
	handle     *pcap.Handle
	ifaceName  string
	outputFile *os.File
	stopChan   chan struct{}
	wg         sync.WaitGroup
	mu         sync.Mutex
}

type PacketStats struct {
	TotalPackets uint64
	OSPFPackets  uint64
	HelloPackets uint64
	LastReceived time.Time
	mu           sync.RWMutex
}

// 创建新的OSPF服务器
func NewOSPFServer(ifaceName, outputPath string) (*OSPFServer, error) {
	// 打开网络接口
	handle, err := pcap.OpenLive(ifaceName, snapLen, promiscuous, timeout)
	if err != nil {
		return nil, fmt.Errorf("打开网络接口失败: %v", err)
	}

	// 设置BPF过滤器，只捕获OSPF数据包
	err = handle.SetBPFFilter("ip proto 89")
	if err != nil {
		handle.Close()
		return nil, fmt.Errorf("设置过滤器失败: %v", err)
	}

	// 打开输出文件
	file, err := os.OpenFile(outputPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		handle.Close()
		return nil, fmt.Errorf("打开输出文件失败: %v", err)
	}

	return &OSPFServer{
		handle:     handle,
		ifaceName:  ifaceName,
		outputFile: file,
		stopChan:   make(chan struct{}),
	}, nil
}

// 启动服务器
func (s *OSPFServer) Start() error {
	logrus.Infof("OSPF服务器启动，监听接口: %s", s.ifaceName)

	s.wg.Add(1)
	go s.receivePackets()

	return nil
}

// 停止服务器
func (s *OSPFServer) Stop() {
	logrus.Infof("正在停止OSPF服务器...")
	close(s.stopChan)
	s.wg.Wait()

	if s.outputFile != nil {
		s.outputFile.Close()
	}
	if s.handle != nil {
		s.handle.Close()
	}
	logrus.Infof("OSPF服务器已停止")
}

// 接收数据包
func (s *OSPFServer) receivePackets() {
	defer s.wg.Done()

	var stats PacketStats
	packetSource := gopacket.NewPacketSource(s.handle, s.handle.LinkType())

	for {
		select {
		case <-s.stopChan:
			return
		default:
			packet, err := packetSource.NextPacket()
			if err != nil {
				continue
			}

			stats.mu.Lock()
			stats.TotalPackets++
			stats.LastReceived = time.Now()
			stats.mu.Unlock()

			// 解析IP层
			ipLayer := packet.Layer(layers.LayerTypeIPv4)
			if ipLayer == nil {
				continue
			}

			ip, _ := ipLayer.(*layers.IPv4)
			if ip.Protocol != layers.IPProtocolOSPF {
				continue
			}

			stats.mu.Lock()
			stats.OSPFPackets++
			stats.mu.Unlock()

			// 获取OSPF数据
			ospfData := packet.ApplicationLayer()
			if ospfData == nil || len(ospfData.Payload()) < OSPFHeaderLen {
				continue
			}

			payload := ospfData.Payload()
			if payload[0] != OSPFVersion || payload[1] != OSPFHelloType {
				continue
			}

			stats.mu.Lock()
			stats.HelloPackets++
			stats.mu.Unlock()

			// 提取Dead Interval字段中的数据
			if len(payload) >= OSPFHeaderLen+HelloPacketLen {
				deadInterval := payload[OSPFHeaderLen+8 : OSPFHeaderLen+12]
				s.writeToFile(ip.SrcIP.String(), deadInterval)
			}

			// 每10秒打印一次统计信息
			if time.Since(stats.LastReceived) >= 10*time.Second {
				stats.mu.RLock()
				logrus.Infof("统计信息 - 总数据包: %d, OSPF包: %d, Hello包: %d", stats.TotalPackets, stats.OSPFPackets, stats.HelloPackets)
				stats.mu.RUnlock()
			}
		}
	}
}

// 写入数据到文件
func (s *OSPFServer) writeToFile(srcIP string, data []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()

	timestamp := time.Now().Format("2006-01-02 15:04:05.000")
	entry := fmt.Sprintf("[%s] 从 %s 接收到数据: %v\n", timestamp, srcIP, data)

	if _, err := s.outputFile.WriteString(entry); err != nil {
		logrus.Errorf("写入文件失败: %v", err)
		return
	}

	s.outputFile.Sync()
}

// 列出所有可用的网络接口
func listInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		logrus.Errorf("查找网络接口失败: %v", err)
		return
	}

	fmt.Println("可用的网络接口:")
	for _, device := range devices {
		logrus.Infof("名称: %s\n", device.Name)
		logrus.Infof("描述: %s\n", device.Description)
		logrus.Infof("IP地址: \n")
		for _, address := range device.Addresses {
			logrus.Infof("  %s\n", address.IP)
		}
		logrus.Println()
	}
}

func main() {
	var (
		ifaceName  = flag.String("i", "", "网络接口名称")
		outputPath = flag.String("o", "ospf_data.txt", "输出文件路径")
		list       = flag.Bool("list", false, "列出所有网络接口")
	)
	flag.Parse()

	if *list {
		listInterfaces()
		return
	}

	if *ifaceName == "" {
		logrus.Fatal("请指定网络接口名称")
	}

	// 创建服务器
	server, err := NewOSPFServer(*ifaceName, *outputPath)
	if err != nil {
		logrus.Fatalf("创建服务器失败: %v", err)
	}

	// 启动服务器
	if err := server.Start(); err != nil {
		logrus.Fatalf("启动服务器失败: %v", err)
	}

	// 处理中断信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// 等待中断信号
	sig := <-sigChan
	logrus.Infof("接收到信号 %v, 正在关闭...", sig)

	// 优雅关闭
	server.Stop()
}
